#include "shape.h"

#include <set>
#include <opencv2/opencv.hpp>

const int Shape::P_ROUND = 1;
const int Shape::P_SQUARE = 0;
const int Shape::MOVE_VERTEX = 0;
const int Shape::NEAR_VERTEX = 1;
const int Shape::PEN_WIDTH = 2;

// Initialize static members
QColor Shape::line_color;
QColor Shape::fill_color;
QColor Shape::select_line_color;
QColor Shape::select_fill_color;
QColor Shape::vertex_fill_color;
QColor Shape::hvertex_fill_color;
int Shape::point_type = Shape::P_ROUND;
int Shape::point_size = 8;
double Shape::scale = 1.0;

// Helper function
bool Shape::isValidShapeType(const QString& value) const
{
	static const std::set<QString> valid_shape_types = {
		"polygon",
		"rectangle",
		"point",
		"line",
		"circle",
		"linestrip",
		"points",
		"mask"
	};
	return valid_shape_types.count(value) > 0;
}

// Constructor implementation
Shape::Shape(
	const QString& label,
	const QColor& line_color,
	const QString& shape_type,
	const std::map<QString, QString>& flags,
	int group_id,
	const QString& description,
	const cv::Mat& mask
)
	: label(label),
	  group_id(group_id),
	  shape_type(shape_type.isEmpty() ? "polygon" : shape_type),
	  fill(false),
	  selected(false),
	  flags(flags),
	  description(description),
	  mask(mask),
	  _highlightIndex(-1),
	  _highlightMode(NEAR_VERTEX),
	  _closed(false)
{
	if (line_color.isValid())
	{
		this->line_color = line_color;
	}

	if (!isValidShapeType(this->shape_type))
	{
		throw std::invalid_argument("Unexpected shape_type: " + this->shape_type.toStdString());
	}

	_highlightSettings[NEAR_VERTEX] = std::make_pair(4.0, P_ROUND);
	_highlightSettings[MOVE_VERTEX] = std::make_pair(1.5, P_SQUARE);
}

QPointF Shape::_scale_point(const QPointF& point) const
{
	return QPointF(point.x() * scale, point.y() * scale);
}

void Shape::setShapeRefined(
	const QString& shape_type,
	const QList<QPointF>& points,
	const QList<int>& point_labels,
	const cv::Mat& mask
)
{
	_shape_raw = std::make_tuple(this->shape_type, this->points, this->point_labels);
	this->shape_type = shape_type;
	this->points = points;
	this->point_labels = point_labels;
	this->mask = mask;
}

void Shape::restoreShapeRaw()
{
	if (std::get<0>(_shape_raw).isEmpty())
	{
		return;
	}
	this->shape_type = std::get<0>(_shape_raw);
	this->points = std::get<1>(_shape_raw);
	this->point_labels = std::get<2>(_shape_raw);
	_shape_raw = std::make_tuple("", QList<QPointF>(), QList<int>());
}

const QString& Shape::getShapeType() const
{
	return shape_type;
}

void Shape::setShapeType(const QString& value)
{
	QString val = value.isEmpty() ? "polygon" : value;
	if (!isValidShapeType(val))
	{
		throw std::invalid_argument("Unexpected shape_type: " + val.toStdString());
	}
	this->shape_type = val;
}

void Shape::close()
{
	_closed = true;
}

void Shape::addPoint(const QPointF& point, int label)
{
	if (!points.empty() && point == points[0])
	{
		close();
	}
	else
	{
		points.push_back(point);
		point_labels.push_back(label);
	}
}

bool Shape::canAddPoint() const
{
	return shape_type == "polygon" || shape_type == "linestrip";
}

QPointF Shape::popPoint()
{
	if (!points.empty())
	{
		if (!point_labels.empty())
		{
			point_labels.pop_back();
		}
		QPointF p = points.back();
		points.pop_back();
		return p;
	}
	return QPointF();
}

void Shape::insertPoint(int i, const QPointF& point, int label)
{
	points.insert(points.begin() + i, point);
	point_labels.insert(point_labels.begin() + i, label);
}

void Shape::removePoint(int i)
{
	if (!canAddPoint())
	{
		qWarning("Cannot remove point from: shape_type=%s", shape_type.toStdString().c_str());
		return;
	}

	if ((shape_type == "polygon" && points.size() <= 3) ||
		(shape_type == "linestrip" && points.size() <= 2))
	{
		qWarning("Cannot remove point from: shape_type=%s, len(points)=%d",
		         shape_type.toStdString().c_str(), static_cast<int>(points.size()));
		return;
	}

	points.erase(points.begin() + i);
	point_labels.erase(point_labels.begin() + i);
}

bool Shape::isClosed() const
{
	return _closed;
}

void Shape::setOpen()
{
	_closed = false;
}

void Shape::paint(QPainter* painter)
{
	if (mask.empty() && points.empty())
	{
		return;
	}

	QColor color = selected ? select_line_color : line_color;
	QPen pen(color);
	pen.setWidth(PEN_WIDTH);
	painter->setPen(pen);

	if (!mask.empty())
	{
		// Step 1: Create image_to_draw with zeros and 4 channels (RGBA)
		cv::Mat imageToDraw(mask.rows, mask.cols, CV_8UC4, cv::Scalar(0, 0, 0, 0));

		// Step 2: Determine fill color based on selection state
		QColor fillColorToUse = selected ? select_fill_color : fill_color;
		int r = fillColorToUse.red();
		int g = fillColorToUse.green();
		int b = fillColorToUse.blue();
		int a = fillColorToUse.alpha();

		// Step 3: Set pixels in image_to_draw where mask is true to fill_color
		imageToDraw.setTo(cv::Scalar(b, g, r, a), mask);

		// Step 4: Convert image_to_draw to QImage
		QImage qimage(imageToDraw.data, imageToDraw.cols, imageToDraw.rows, imageToDraw.step, QImage::Format_RGBA8888);

		// Ensure the data is not freed after QImage is destroyed
		qimage.bits();

		// Step 5: Scale the QImage
		int newWidth = static_cast<int>(qimage.width() * scale);
		int newHeight = static_cast<int>(qimage.height() * scale);
		qimage = qimage.scaled(newWidth, newHeight, Qt::IgnoreAspectRatio, Qt::SmoothTransformation);

		// Step 6: Draw the image at the scaled starting point
		painter->drawImage(_scale_point(points[0]), qimage);

		// Step 7: Create a QPainterPath for the contours
		QPainterPath linePath;

		// Step 8: Pad the mask and find contours
		cv::Mat paddedMask;
		cv::copyMakeBorder(mask, paddedMask, 1, 1, 1, 1, cv::BORDER_CONSTANT, cv::Scalar(0));
		std::vector<std::vector<cv::Point>> contours;
		cv::findContours(paddedMask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);

		// Step 9 and 10: Adjust contour coordinates and add to path
		for (const auto& contour : contours)
		{
			if (contour.empty())
				continue;

			// Move to starting point
			QPointF startPoint(contour[0].x + points[0].x(), contour[0].y + points[0].y());
			startPoint = _scale_point(startPoint);
			linePath.moveTo(startPoint);

			// Add lines to path
			for (size_t i = 1; i < contour.size(); ++i)
			{
				QPointF pt(contour[i].x + points[0].x(), contour[i].y + points[0].y());
				pt = _scale_point(pt);
				linePath.lineTo(pt);
			}
		}

		// Step 11: Draw the path
		painter->drawPath(linePath);
	}

	if (!points.empty())
	{
		QPainterPath line_path;
		QPainterPath vrtx_path;
		QPainterPath negative_vrtx_path;

		if (shape_type == "rectangle" || shape_type == "mask")
		{
			assert(points.size() == 1 || points.size() == 2);
			if (points.size() == 2)
			{
				QRectF rectangle(_scale_point(points[0]), _scale_point(points[1]));
				line_path.addRect(rectangle);
			}
			if (shape_type == "rectangle")
			{
				for (size_t i = 0; i < points.size(); ++i)
				{
					drawVertex(vrtx_path, i);
				}
			}
		}
		else if (shape_type == "circle")
		{
			assert(points.size() == 1 || points.size() == 2);
			if (points.size() == 2)
			{
				double radius = std::hypot(
					_scale_point(points[0]).x() - _scale_point(points[1]).x(),
					_scale_point(points[0]).y() - _scale_point(points[1]).y()
				);
				line_path.addEllipse(_scale_point(points[0]), radius, radius);
			}
			for (size_t i = 0; i < points.size(); ++i)
			{
				drawVertex(vrtx_path, i);
			}
		}
		else if (shape_type == "linestrip")
		{
			line_path.moveTo(_scale_point(points[0]));
			for (size_t i = 0; i < points.size(); ++i)
			{
				line_path.lineTo(_scale_point(points[i]));
				drawVertex(vrtx_path, i);
			}
		}
		else if (shape_type == "points")
		{
			assert(points.size() == point_labels.size());
			for (size_t i = 0; i < point_labels.size(); ++i)
			{
				int point_label = point_labels[i];
				if (point_label == 1)
				{
					drawVertex(vrtx_path, i);
				}
				else
				{
					drawVertex(negative_vrtx_path, i);
				}
			}
		}
		else
		{
			line_path.moveTo(_scale_point(points[0]));
			for (size_t i = 0; i < points.size(); ++i)
			{
				line_path.lineTo(_scale_point(points[i]));
				drawVertex(vrtx_path, i);
			}
			if (isClosed())
			{
				line_path.lineTo(_scale_point(points[0]));
			}
		}

		painter->drawPath(line_path);
		if (!vrtx_path.isEmpty())
		{
			painter->drawPath(vrtx_path);
			painter->fillPath(vrtx_path, _vertex_fill_color);
		}
		if (fill && mask.empty())
		{
			QColor color = selected ? select_fill_color : fill_color;
			painter->fillPath(line_path, color);
		}

		pen.setColor(QColor(255, 0, 0, 255));
		painter->setPen(pen);
		painter->drawPath(negative_vrtx_path);
		painter->fillPath(negative_vrtx_path, QColor(255, 0, 0, 255));
	}
}

void Shape::drawVertex(QPainterPath& path, int i)
{
	double d = point_size;
	int shape = point_type;
	QPointF point = _scale_point(points[i]);
	if (i == _highlightIndex)
	{
		double size = _highlightSettings.at(_highlightMode).first;
		shape = _highlightSettings.at(_highlightMode).second;
		d *= size;
	}
	_vertex_fill_color = (_highlightIndex != -1) ? hvertex_fill_color : vertex_fill_color;

	if (shape == P_SQUARE)
	{
		path.addRect(point.x() - d / 2, point.y() - d / 2, d, d);
	}
	else if (shape == P_ROUND)
	{
		path.addEllipse(point, d / 2.0, d / 2.0);
	}
	else
	{
		assert(false && "unsupported vertex shape");
	}
}

int Shape::nearestVertex(const QPointF& point, double epsilon) const
{
	double min_distance = std::numeric_limits<double>::infinity();
	int min_i = -1;
	QPointF scaled_point(point.x() * scale, point.y() * scale);
	for (size_t i = 0; i < points.size(); ++i)
	{
		QPointF p(points[i].x() * scale, points[i].y() * scale);
		double dist = std::hypot(p.x() - scaled_point.x(), p.y() - scaled_point.y());
		if (dist <= epsilon && dist < min_distance)
		{
			min_distance = dist;
			min_i = static_cast<int>(i);
		}
	}
	return min_i;
}

int Shape::nearestEdge(const QPointF& point, double epsilon) const
{
	double min_distance = std::numeric_limits<double>::infinity();
	int post_i = -1;
	QPointF scaled_point(point.x() * scale, point.y() * scale);
	for (size_t i = 0; i < points.size(); ++i)
	{
		QPointF start = points[(i - 1 + points.size()) % points.size()];
		QPointF end = points[i];
		start = QPointF(start.x() * scale, start.y() * scale);
		end = QPointF(end.x() * scale, end.y() * scale);

		// Calculate distance from point to line segment
		double x0 = scaled_point.x(), y0 = scaled_point.y();
		double x1 = start.x(), y1 = start.y();
		double x2 = end.x(), y2 = end.y();
		double num = std::abs((x2 - x1) * (y1 - y0) - (x1 - x0) * (y2 - y1));
		double den = std::hypot(x2 - x1, y2 - y1);
		double dist = num / den;

		if (dist <= epsilon && dist < min_distance)
		{
			min_distance = dist;
			post_i = static_cast<int>(i);
		}
	}
	return post_i;
}

bool Shape::containsPoint(const QPointF& point) const
{
	if (!mask.empty())
	{
		// Calculate the relative coordinates
		int y = static_cast<int>(std::round(point.y() - points[0].y()));
		int x = static_cast<int>(std::round(point.x() - points[0].x()));

		// Clamp the coordinates to the mask boundaries
		y = std::max(0, std::min(y, mask.rows - 1));
		x = std::max(0, std::min(x, mask.cols - 1));

		// Check if the mask at (y, x) is non-zero
		return mask.at<uchar>(y, x) != 0;
	}
	else
	{
		// Fallback to checking if the point is within the path
		return makePath().contains(point);
	}
}

QPainterPath Shape::makePath() const
{
	QPainterPath path;
	if (shape_type == "rectangle" || shape_type == "mask")
	{
		if (points.size() == 2)
		{
			path.addRect(QRectF(points[0], points[1]));
		}
	}
	else if (shape_type == "circle")
	{
		if (points.size() == 2)
		{
			double radius = std::hypot(points[0].x() - points[1].x(), points[0].y() - points[1].y());
			path.addEllipse(points[0], radius, radius);
		}
	}
	else if (!points.empty())
	{
		path.moveTo(points[0]);
		for (size_t i = 1; i < points.size(); ++i)
		{
			path.lineTo(points[i]);
		}
	}
	return path;
}

QRectF Shape::boundingRect() const
{
	return makePath().boundingRect();
}

void Shape::moveBy(const QPointF& offset)
{
	for (auto& point : points)
	{
		point += offset;
	}
}

void Shape::moveVertexBy(int i, const QPointF& offset)
{
	points[i] += offset;
}

void Shape::highlightVertex(int i, int action)
{
	_highlightIndex = i;
	_highlightMode = action;
}

void Shape::highlightClear()
{
	_highlightIndex = -1;
}

Shape* Shape::copy() const
{
	return new Shape(*this);
}

int Shape::size() const
{
	return static_cast<int>(points.size());
}

QPointF Shape::operator[](int key) const
{
	return points[key];
}

void Shape::setPoint(int key, const QPointF& value)
{
	points[key] = value;
}
